/*
Copyright (c) 2024 Seldon Technologies Ltd.

Use of this software is governed by
(1) the license included in the LICENSE file or
(2) if the license included in the LICENSE file is the Business Source License 1.1,
the Change License after the Change Date as each is defined in accordance with the LICENSE file.
*/

package pipeline

import (
	"context"
	"fmt"
	"math"
	"os"
	"sync/atomic"
	"time"

	"github.com/cenkalti/backoff/v4"
	grpc_retry "github.com/grpc-ecosystem/go-grpc-middleware/retry"
	"github.com/sirupsen/logrus"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/protobuf/encoding/protojson"

	chainer "github.com/seldonio/seldon-core/apis/go/v2/mlops/chainer"
	"github.com/seldonio/seldon-core/apis/go/v2/mlops/scheduler"
	pb "github.com/seldonio/seldon-core/apis/go/v2/mlops/scheduler"
	"github.com/seldonio/seldon-core/components/tls/v2/pkg/tls"

	"github.com/seldonio/seldon-core/scheduler/v2/pkg/kafka/pipeline/status"
	"github.com/seldonio/seldon-core/scheduler/v2/pkg/store/pipeline"
	"github.com/seldonio/seldon-core/scheduler/v2/pkg/util"
)

const (
	SubscriberName = "seldon-pipelinegateway"
)

type PipelineSchedulerClient struct {
	logger                logrus.FieldLogger
	conn                  *grpc.ClientConn
	callOptions           []grpc.CallOption
	pipelineStatusUpdater status.PipelineStatusUpdater
	pipelineInferer       PipelineInferer
	stop                  atomic.Bool
	ready                 atomic.Bool
	tlsOptions            *tls.TLSOptions
}

func NewPipelineSchedulerClient(logger logrus.FieldLogger, pipelineStatusUpdater status.PipelineStatusUpdater, pipelineInferer PipelineInferer, tlsOptions *tls.TLSOptions) *PipelineSchedulerClient {
	opts := []grpc.CallOption{
		grpc.MaxCallSendMsgSize(math.MaxInt32),
		grpc.MaxCallRecvMsgSize(math.MaxInt32),
	}

	return &PipelineSchedulerClient{
		logger:                logger.WithField("source", "PipelineSchedulerClient"),
		callOptions:           opts,
		pipelineStatusUpdater: pipelineStatusUpdater,
		pipelineInferer:       pipelineInferer,
		tlsOptions:            tlsOptions,
	}
}

func (pc *PipelineSchedulerClient) IsConnected() bool {
	return pc.ready.Load()
}

func (pc *PipelineSchedulerClient) connectToScheduler(host string, plainTxtPort int, tlsPort int) error {
	logger := pc.logger.WithField("func", "ConnectToScheduler")
	var err error
	if pc.conn != nil {
		err = pc.conn.Close()
		if err != nil {
			logger.WithError(err).Error("Failed to close previous grpc connection to scheduler")
		}
	}

	var transCreds credentials.TransportCredentials
	var port int
	if pc.tlsOptions.Cert == nil {
		logger.Info("Starting plaintxt client to scheduler")
		transCreds = insecure.NewCredentials()
		port = plainTxtPort
	} else {
		logger.Info("Starting TLS client to scheduler")
		transCreds = pc.tlsOptions.Cert.CreateClientTransportCredentials()
		port = tlsPort
	}

	kacp := util.GetClientKeepAliveParameters()

	// note: retry is done in the caller
	opts := []grpc.DialOption{
		grpc.WithTransportCredentials(transCreds),
		grpc.WithKeepaliveParams(kacp),
	}
	logger.Infof("Connecting to scheduler at %s:%d", host, port)
	conn, err := grpc.NewClient(fmt.Sprintf("%s:%d", host, port), opts...)
	if err != nil {
		return err
	}
	pc.conn = conn
	return nil
}

func (pc *PipelineSchedulerClient) Stop() {
	pc.stop.Store(true)
	if pc.conn != nil {
		_ = pc.conn.Close()
	}
}

func (pc *PipelineSchedulerClient) Start(host string, plainTxtPort int, tlsPort int) error {
	logger := pc.logger.WithField("func", "Start")
	for {
		if pc.stop.Load() {
			logger.Info("Stopping")
			return nil
		}
		err := pc.connectToScheduler(host, plainTxtPort, tlsPort)
		if err != nil {
			logger.WithError(err).Fatalf("Failed to connect to scheduler")
		}
		logger := pc.logger.WithField("func", "Start")
		logFailure := func(err error, delay time.Duration) {
			logger.WithError(err).Errorf("Scheduler not ready")
		}
		backOffExp := util.GetClientExponentialBackoff()
		err = backoff.RetryNotify(pc.SubscribePipelineEvents, backOffExp, logFailure)
		if err != nil {
			logger.WithError(err).Fatal("Failed to start pipeline gateway client")
			return err
		}
		logger.Info("Subscribe ended")
	}
}

func getSubscriberName() string {
	podName := os.Getenv("POD_NAME")
	if podName == "" {
		return SubscriberName
	}
	return podName
}

func getSubsriberIp() (string, error) {
	podIp := os.Getenv("POD_IP")
	if podIp == "" {
		return "", fmt.Errorf("POD_IP environment variable is not set")
	}
	return podIp, nil
}

func (pc *PipelineSchedulerClient) SubscribePipelineEvents() error {
	logger := pc.logger.WithField("func", "SubscribePipelineEvents")

	stream, processor, closeStream, err := pc.setupSubscription(logger)
	if err != nil {
		return err
	}

	defer pc.cleanup(closeStream)
	return pc.processEventStream(stream, processor, logger)
}

func (pc *PipelineSchedulerClient) setupSubscription(logger *logrus.Entry) (scheduler.Scheduler_SubscribePipelineStatusClient, *EventProcessor, context.CancelFunc, error) {
	grpcClient := scheduler.NewSchedulerClient(pc.conn)

	subscriberName := getSubscriberName()
	subscriberIp, err := getSubsriberIp()
	if err != nil {
		return nil, nil, nil, err
	}

	logger.Infof("Subscriber (%s, %s) subscribing to pipeline status events", subscriberName, subscriberIp)
	ctx, cancel := context.WithCancel(context.Background())
	stream, err := grpcClient.SubscribePipelineStatus(
		ctx,
		&scheduler.PipelineSubscriptionRequest{
			SubscriberName:    subscriberName,
			SubscriberIp:      subscriberIp,
			IsPipelineGateway: true,
		},
		grpc_retry.WithMax(util.MaxGRPCRetriesOnStream),
	)

	if err != nil {
		cancel()
		return nil, nil, nil, err
	}

	processor := &EventProcessor{
		client:         pc,
		grpcClient:     grpcClient,
		subscriberName: subscriberName,
		logger:         logger,
	}

	pc.ready.Store(true)
	return stream, processor, cancel, nil
}

func (pc *PipelineSchedulerClient) processEventStream(stream scheduler.Scheduler_SubscribePipelineStatusClient, processor *EventProcessor, logger *logrus.Entry) error {
	for {
		if pc.stop.Load() {
			logger.Info("Stopping")
			return nil
		}

		event, err := stream.Recv()
		if err != nil {
			logger.WithError(err).Error("event recv failed")
			return err
		}

		processor.HandleEvent(event)
	}
}

func (pc *PipelineSchedulerClient) cleanup(closeStream context.CancelFunc) {
	pc.ready.Store(false)
	pc.logger.Info("Closing connection to scheduler")
	closeStream()
}

type EventProcessor struct {
	client         *PipelineSchedulerClient
	grpcClient     scheduler.SchedulerClient
	subscriberName string
	logger         *logrus.Entry
}

func NewEventProcessor(
	client *PipelineSchedulerClient,
	grpcClient scheduler.SchedulerClient,
	subscriberName string,
	logger logrus.FieldLogger,
) *EventProcessor {
	return &EventProcessor{
		client:         client,
		grpcClient:     grpcClient,
		subscriberName: subscriberName,
		logger:         logger.WithField("component", "EventProcessor"),
	}
}

func (ep *EventProcessor) HandleEvent(event *scheduler.PipelineStatusResponse) {
	// The expected contract is just the latest version will be sent to us
	if len(event.Versions) != 1 {
		ep.logger.Info("Expected a single model version", "numVersions", len(event.Versions), "name", event.GetPipelineName())
		return
	}

	switch event.Operation {
	case pb.PipelineStatusResponse_PipelineDelete:
		ep.handleDeletePipeline(event)
	case pb.PipelineStatusResponse_PipelineCreate:
		ep.handleCreateOrUpdatePipeline(event)
	}
}

func (ep *EventProcessor) handleDeletePipeline(event *scheduler.PipelineStatusResponse) {
	pv, err := pipeline.CreatePipelineVersionWithStateFromProto(event.Versions[0])
	if err != nil {
		ep.reportFailure(
			chainer.PipelineUpdateMessage_Delete,
			nil,
			fmt.Sprintf("No existing pipeline %s to delete", event.PipelineName),
			event.Timestamp,
			nil,
		)
		return
	}

	err = ep.client.pipelineInferer.DeletePipeline(event.PipelineName, false)
	if err != nil {
		ep.reportFailure(
			chainer.PipelineUpdateMessage_Delete,
			pv,
			fmt.Sprintf("Failed to delete pipeline %s", event.PipelineName),
			event.Timestamp,
			err,
		)
		return
	}

	message := fmt.Sprintf("Pipeline %s deleted", event.PipelineName)
	ep.reportSuccess(chainer.PipelineUpdateMessage_Delete, pv, message, event.Timestamp)
}

func (ep *EventProcessor) handleCreateOrUpdatePipeline(event *scheduler.PipelineStatusResponse) {
	pv, err := pipeline.CreatePipelineVersionWithStateFromProto(event.Versions[0])
	if err != nil {
		ep.reportFailure(
			chainer.PipelineUpdateMessage_Create,
			nil,
			fmt.Sprintf("Failed to create pipeline version for pipeline %s with %s", event.PipelineName, protojson.Format(event)),
			event.Timestamp,
			err,
		)
		return
	}

	ep.logger.Debugf("Processing pipeline %s version %d with state %s", pv.Name, pv.Version, pv.State.Status.String())
	ep.client.pipelineStatusUpdater.Update(pv)

	_, err = ep.client.pipelineInferer.LoadOrStorePipeline(pv.Name, false, false)
	if err != nil {
		ep.reportFailure(
			chainer.PipelineUpdateMessage_Create,
			pv,
			fmt.Sprintf("Failed to load/store pipeline %s", pv.Name),
			event.Timestamp,
			err,
		)
		return
	}

	message := fmt.Sprintf("Pipeline %s loaded", event.PipelineName)
	ep.reportSuccess(chainer.PipelineUpdateMessage_Create, pv, message, event.Timestamp)
}

func (ep *EventProcessor) reportSuccess(op chainer.PipelineUpdateMessage_PipelineOperation, pv *pipeline.PipelineVersion, message string, timestamp uint64) {
	ep.logger.Info(message)
	go ep.sendPipelineStatusEvent(op, pv, true, message, timestamp)
}

func (ep *EventProcessor) reportFailure(op chainer.PipelineUpdateMessage_PipelineOperation, pv *pipeline.PipelineVersion, message string, timestamp uint64, err error) {
	if err != nil {
		ep.logger.WithError(err).Error(message)
	} else {
		ep.logger.Error(message)
	}

	if pv != nil {
		go ep.sendPipelineStatusEvent(op, pv, false, message, timestamp)
	}
}

func (ep *EventProcessor) sendPipelineStatusEvent(
	op chainer.PipelineUpdateMessage_PipelineOperation,
	pv *pipeline.PipelineVersion,
	success bool,
	reason string,
	timestamp uint64,
) {
	callOpts := []grpc.CallOption{
		grpc_retry.WithMax(5), // retry up to 5 times
		grpc_retry.WithBackoff(grpc_retry.BackoffExponentialWithJitter(100*time.Millisecond, 2.0)),
	}

	_, err := ep.grpcClient.PipelineStatusEvent(
		context.Background(),
		&chainer.PipelineUpdateStatusMessage{
			Update: &chainer.PipelineUpdateMessage{
				Op:        op,
				Pipeline:  pv.Name,
				Version:   pv.Version,
				Uid:       pv.UID,
				Stream:    ep.subscriberName,
				Timestamp: timestamp,
			},
			Success: success,
			Reason:  reason,
		},
		callOpts...,
	)
	if err != nil {
		ep.logger.WithError(err).Errorf("Failed to send pipeline status event for pipeline %s", pv.Name)
	}
}
